from langchain_community.chat_models import ChatZhipuAI
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_chroma import Chroma

from langchain_core.runnables import RunnableLambda, RunnablePassthrough

text_loader = TextLoader("../data/背影.txt", encoding="utf-8")

documents = text_loader.load()

embedding = ZhipuAIEmbeddings()

# vector_db = Chroma.from_documents(
#     documents=documents,
#     embedding=embedding,
#     persist_directory="../data/chroma_db"
# )

vector_db = Chroma(
    persist_directory="../data/chroma_db",
    embedding_function=embedding
)

retriever = RunnableLambda(vector_db.similarity_search).bind(k=1)

model = ChatZhipuAI(
    model="glm-4-plus",
    temprature=0.5
)

prompts = ChatPromptTemplate.from_messages([
    ("system", "请根据上下文回答问题"),
    ("human", "问题:{question} 上下文:{context}")
])

chain = {"question":RunnablePassthrough(),"context":retriever} | prompts | model

response = chain.invoke("父亲买了什么")
print(response.content)